"""
角色：规划、执行、收集、团队
链式任务规划----2024年法国巴黎奥运会女子10米跳水冠军的家乡在哪里
0.定义图节点状态对象
1.初始planner/executor/gatherer agent ,需要大模型
2.定义节点 _planner_node/_executor_node/_gatherer_node
3.路由定义 _router_edge ，当执行任务数量和收集信息数量相等时 转到 _gatherer_node节点，否则装_executor_node节点继续执行
4.构建和编译graph 节点和团队agent绑定，边的构建
5.执行graph.invoke
"""

from model_utils import getLLM
from execute import Executor
from gather import Gatherer
from plan import Planner
from typing import Annotated,TypedDict
from operator import add

_llm = getLLM()
_planner_agent = Planner(_llm)
_executor_agent = Executor(_llm)
_gatherer_agent = Gatherer(_llm)

class PlanState(TypedDict):
    query:str
    task_list:list[str]
    infos:Annotated[list[str],add]
    result:str

def _planner_node(state):
    return _planner_agent(state)

# 结束使用条件边控制
def _executor_node(state):
    # 取当前完成的任务得到信息
    _infos = state.get('infos',[])
    _task_index = len(_infos)
    # 取当前待执行的任务
    _task_name = state["task_list"][_task_index]

    _rt = _executor_agent({"task_name":_task_name,"infos":_infos})
    return {"infos":[_rt]}

def _gatherer_node(state):
    _rt = _gatherer_agent({"query":state['query'],"infos":state['infos']})
    return {"result":_rt}

def _router_edge(state):
    _task_list_len = len(state["task_list"])
    _infos_len = len(state["infos"])

    if(_task_list_len==_infos_len):
        return "_gatherer_node"
    else:
        return "_executor_node"

from langgraph.graph import StateGraph,START,END,MessagesState,MessageGraph
_builder = StateGraph(PlanState)
_builder.add_node("_planner_node",_planner_node)
_builder.add_node("_executor_node",_executor_node)
_builder.add_node("_gatherer_node",_gatherer_node)

_builder.add_edge(START,"_planner_node")
_builder.add_edge("_planner_node","_executor_node")
_builder.add_conditional_edges("_executor_node",_router_edge,["_executor_node","_gatherer_node"])
# _builder.add_conditional_edges("_executor_node",router_edge,{"_gatherer_node":"_gatherer_node","_executor_node":"_executor_node"})
_builder.add_edge("_gatherer_node",END)

_graph = _builder.compile()

_query = "2024年法国巴黎奥运会男子篮球美国队的队长正在效力的NBA球队是哪一支"

_rt = _graph.invoke({"query":_query})
print(_rt)
print("结果：",_rt.get("result",""))

_graph.get_graph().draw_mermaid_png(output_file_path="team.png")